from monai.networks.nets.basic_unet import BasicUNet
import torch.nn as nn 
import torch 
from einops import rearrange

def infer2d(image3d, network):
    # image3d: batch, channels, D, W, H
    index = 0
    res_pred = None
    for y in torch.split(image3d, 1, dim=2):
        if index == 0:
            y = y.squeeze(2)
            pred_2d = network(y)
            res_pred = pred_2d.unsqueeze(2)
        else:
            y = y.squeeze(2)
            pred_2d = network(y)
            res_pred = torch.cat((res_pred, pred_2d.unsqueeze(2)), dim=2)

        index += 1
    return res_pred

class UNet2D(nn.Module):
    def __init__(self, in_ch, out_ch) -> None:
        super().__init__()
        self.model = BasicUNet(spatial_dims=2, 
                                in_channels=in_ch, 
                                out_channels=out_ch,
                                features=[16, 16, 32, 64, 128, 16])

        self.out_ch = out_ch

    def forward(self, x):
        # x: (b, c, d, w, h)
        # b, c, d, w, h = x.shape
        # x = rearrange(x, "b c d w h -> (b d) c w h")
        # x = self.model(x)
        # x = rearrange(x, "(b d) c w h -> b c d w h", b=b, c=self.out_ch)

        x = infer2d(x, self.model)
        
        return x 
        


